from typing import Union, Dict

import torch
import numpy as np


class ToTensor:
    def __init__(self, keys=None) -> None:
        self.keys = keys
    
    def _apply(self, sample: np.ndarray) -> torch.Tensor:
        """
        Convert a numpy array to a PyTorch tensor.
        
        :param sample: numpy array.
        :return: PyTorch tensor.
        """
        return torch.from_numpy(sample)

    def __call__(self, data: Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]) -> Union[torch.Tensor, Dict[str, Union[torch.Tensor, float]]]:
        if isinstance(data, dict):
            for key in self.keys:
                if key in data:
                    data[key] = self._apply(data[key])
        elif isinstance(data, np.ndarray):
            data = self._apply(data)
        return data